/*
 * MIT License
 *
 * Copyright (c) 2020 wen.gu <454727014@qq.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

 /***************************************************************************
 * Name: tiny_encoder.c
 *
 * Purpose: implementation a encoder with tiny compress
 *
 * Developer:
 *   wen.gu , 2022-03-04
 *
 * TODO:
 *
 ***************************************************************************/

 /******************************************************************************
 **    INCLUDES
 ******************************************************************************/

#include "tiny_encoder.h"

#include <stdio.h>
#include <string.h>
#include <limits.h>
#include <vector>

#include "lz4/lz4.h"
#include "lz4/lz4hc.h"
#include "zstd/zstd.h"
/******************************************************************************
 **    MACROS
 ******************************************************************************/
#define LOGE printf
#define LOGD printf

#define Y_BLOCK_WIDTH 16
#define UV_BLOCK_WIDTH 8

#define Y_BLOCK_SIZE (Y_BLOCK_WIDTH * Y_BLOCK_WIDTH)
#define UV_BLOCK_SIZE (UV_BLOCK_WIDTH * UV_BLOCK_WIDTH)
/******************************************************************************
 **    VARIABLE DEFINITIONS
 ******************************************************************************/
struct BlockCoordinate
{
    uint32_t x;
    uint32_t y;
};

using BlockData = uint8_t[Y_BLOCK_SIZE];

//using BlockArray = std::vector<BlockCoordinate>;


/******************************************************************************
 **    inner FUNCTION DEFINITIONS
 ******************************************************************************/
#define RESIDUAL_STEP 2

static void block_search_nearest3(const uint8_t* reference_frame, const uint8_t* cur_frame, uint32_t width, uint32_t height, uint32_t block_width, uint32_t& pos_x, uint32_t& pos_y)
{
#if 1
    /**
     * @brief search rule: 
     *  #    #    #             #    #
     *  #    #    #             #    #
     *  #    #    #(cur block)  #    #
     *  #    #    #             #    # 
     *  #    #    #             #    # 
     * total compare block count: 25
     */
#define SEARCH_AREA (5 * block_width)

    uint32_t start_x = (pos_x > SEARCH_AREA) ? (pos_x - SEARCH_AREA) : 0;
    uint32_t start_y = (pos_y > SEARCH_AREA) ? (pos_y - SEARCH_AREA) : 0;

    uint32_t end_x = pos_x + SEARCH_AREA;
    uint32_t end_y = pos_y + SEARCH_AREA;

    end_x = (end_x > width) ? width : end_x;
    end_y = (end_y > height) ? height : end_y;
#else
    // full search
    uint32_t start_x = 0;
    uint32_t start_y = 0;

    uint32_t end_x = width;
    uint32_t end_y = height; 
#endif
    uint32_t out_x = 0;
    uint32_t out_y = 0;
    unsigned int min_diff = UINT_MAX;

    for (size_t i = start_y; i < end_y; i += block_width)
    {
        for (size_t j = start_x; j < end_x; j += block_width)
        {
            unsigned int diff_sum = 0;
            //const uint8_t* cur_buf = cur_frame + i * width + j;
            //const uint8_t* ref_buf = reference_frame + i * width + j;

            for (size_t block_h = 0; block_h < block_width; block_h++)
            {
                for (size_t block_w = 0; block_w < block_width; block_w++)
                {
                    unsigned int c1 = cur_frame[(pos_y + block_h) * width + pos_x + block_w];
                    unsigned int c2 = reference_frame[(i + block_h) * width + j + block_w];
                    diff_sum += (c1 - c2) * (c1 - c2);
                }                
            }

            if (min_diff > diff_sum)
            {
                out_x = j;
                out_y = i;
                min_diff = diff_sum;
            }
        }        
    }
#if 0
    if (min_diff != 0)
    {
        LOGD("[%s.%d]===>: beast diff: %u\n", __FUNCTION__, __LINE__, min_diff);
    }

    compare_block(reference_frame, cur_frame, out_x, out_y, block_width, width);
#endif
    pos_x = out_x;
    pos_y = out_y;
}

static void generate_predictive_frame(const uint8_t* ref_frame, const uint8_t* cur_frame, uint8_t* predictive_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    for (uint32_t i = 0; i < height; i += block_width)
    {
        for (uint32_t j = 0; j < width; j += block_width)
        {
            
            uint32_t pos_x = j;
            uint32_t pos_y = i;
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

            block_search_nearest3(ref_frame, cur_frame, width, height, block_width, pos_x, pos_y);
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            //printf("[%s.%d]===>: cur (%d,%d), nearest(%d,%d)\n", __FUNCTION__, __LINE__, j,i, pos_x, pos_y);

            for (uint32_t block_h = 0; block_h < block_width; block_h++)
            {
                memcpy(predictive_frame + (i + block_h) * width + j, ref_frame + (pos_y + block_h) * width + pos_x, block_width);
            }
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            
            //save_block_data2(predictive_frame, ref_frame, pos_x, pos_y, width, block_width);
        }
    }
}

static void partition_encode(const uint8_t* ref_frame, const uint8_t* enc_frame, uint16_t* coordinate_buf, uint8_t* residual_frame, uint8_t* next_ref_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    for (uint32_t i = 0; i < height; i += block_width)
    {
        for (uint32_t j = 0; j < width; j += block_width)
        {
            
            uint32_t pos_x = j;
            uint32_t pos_y = i;
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

            block_search_nearest3(ref_frame, enc_frame, width, height, block_width, pos_x, pos_y);
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            //printf("[%s.%d]===>: cur (%d,%d), nearest(%d,%d)\n", __FUNCTION__, __LINE__, j,i, pos_x, pos_y);
            *coordinate_buf++ = pos_x;
            *coordinate_buf++ = pos_y;
#if 0
            for (uint32_t block_h = 0; block_h < block_width; block_h++)
            {
                memcpy(predictive_frame + (i + block_h) * width + j, ref_frame + (pos_y + block_h) * width + pos_x, block_width);
                for (uint32_t block_w = 0; block_w < block_width; block_w++)
                {
                    /* code */
                }
                
            }
#else
            for (uint32_t block_h = 0; block_h < block_width; block_h++)
            {
                for (uint32_t block_w = 0; block_w < block_width; block_w++)
                {
                    uint32_t predictive_idx = (i + block_h) * width + j + block_w;
                    uint32_t ref_idx = (pos_y + block_h) * width + pos_x + block_w;
                    //predictive_frame[predictive_idx] = ref_frame[ref_idx];
                    /** calculate residual */
                    //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
                    residual_frame[predictive_idx] = (enc_frame[predictive_idx] - ref_frame[ref_idx]) / RESIDUAL_STEP; // + 0x80;
                    //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
                    next_ref_frame[predictive_idx] = ref_frame[ref_idx] + residual_frame[predictive_idx] * RESIDUAL_STEP;
                    //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
                }
            }
#endif            
        }
    }
}


static void dump_frame(const char* out_file, const uint8_t* frame_buf, uint32_t frame_size)
{
    FILE* fp = fopen(out_file, "wb");
    if (!fp)
    {
        LOGE("open file(%s) failed(%s)\n", out_file, strerror(errno));
        return ;
    }

    int res = fwrite(frame_buf, 1, frame_size, fp);

    fflush(fp);
    fclose(fp);
}
/******************************************************************************
 **    FUNCTION DEFINITIONS
 ******************************************************************************/
TinyEncoder::TinyEncoder()
{

}

TinyEncoder::~TinyEncoder()
{
    for (size_t i = 0; i < TINY_ENC_REF_BUF_CNT; i++)
    {
        if (reference_frame_buf_[i])
        {
            delete[] reference_frame_buf_[i];
        }
    }

    if (out_frame_buf_)
    {
        delete[] out_frame_buf_;
    }
}



bool TinyEncoder::encode(const uint8_t* src, uint32_t src_size, KfaBuffer& out_buf)
{
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    if (!reference_frame_buf_[0])
    {
        //mean first frame, the only 'I' frame
        for (size_t i = 0; i < TINY_ENC_REF_BUF_CNT; i++)
        {
            reference_frame_buf_[i] = new uint8_t[src_size];
        }        
        
        memcpy(reference_frame_buf_[reference_buf_idx_], src, src_size);
        //reference_y_size_ = src_size;
        uint32_t out_buf_size = LZ4_COMPRESSBOUND(src_size);
        //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
        out_buf.buf = new uint8_t[out_buf_size];
        out_buf.release = [](KfaBuffer& buf) 
        {
            if (buf.buf)
            {
                delete[] buf.buf;
            }
        };

        int res_size = LZ4_compress_HC((const char*)src, (char*)out_buf.buf, src_size, out_buf_size, LZ4HC_CLEVEL_MAX);

        if(res_size <= 0)
        {
            return false;
        } 
        //printf("[%s.%d]===>: res_size:\n", __FUNCTION__, __LINE__);
        out_buf.size = res_size;
        return true;
    }

    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    uint32_t width = 1920; //todo fix me??
    uint32_t height = 720;
    uint32_t y_size = width * height;
    uint32_t uv_size = y_size / 2;
    uint32_t frame_size = y_size + uv_size; // only for nv12

    uint32_t y_block_coordinate_size = y_size / Y_BLOCK_SIZE * 4; /** use 2 byte indicat coordinate, (x,y) total 4 byte */
    uint32_t uv_block_coordinate_size = uv_size / UV_BLOCK_SIZE * 4; /** use 2 byte indicat coordinate, (x,y) total 4 byte */
    uint32_t block_coordinate_size = y_block_coordinate_size + uv_block_coordinate_size;
    uint32_t before_size = src_size + block_coordinate_size;
    
    if (!out_frame_buf_)
    {
        out_frame_buf_ = new uint8_t[before_size];
    }
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    uint8_t* out_frame = out_frame_buf_ + block_coordinate_size;
    uint16_t* y_coordinate_buf = (uint16_t*)out_frame_buf_;
    uint16_t* uv_coordinate_buf = (uint16_t*)(out_frame_buf_ + y_block_coordinate_size);
#if 0    

    for (size_t i = 0; i < height; i += Y_BLOCK_WIDTH)
    {
        for (size_t j = 0; j < width; j += Y_BLOCK_WIDTH)
        {
            BlockData block;
            //printf("[%s.%d]===>:y: %d, x:%d\n", __FUNCTION__, __LINE__, i, j);
            frame_get_block(src, i, j, width, block);
            //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            BlockCoordinate pos;
            pos.x = j;
            pos.y = i;
            block_search_nearest(reference_frame_buf_, block, width, height, pos);
            BlockData ref_block;
            //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            frame_get_block(reference_frame_buf_, pos.x, pos.y, width, ref_block);
            //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            //out_blocks.push_back(pos);
            *((uint16_t*)coordinate_buf) = pos.x;
            coordinate_buf += sizeof(uint16_t);
            *((uint16_t*)coordinate_buf) = pos.y;
            coordinate_buf += sizeof(uint16_t);
            output_residual_data(out_frame, block, ref_block, j,i, width);
            //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);            
        }        
    }


#else
#if 1
    uint8_t* ref_buf = reference_frame_buf_[reference_buf_idx_ % TINY_ENC_REF_BUF_CNT];
    uint8_t* next_ref_buf = reference_frame_buf_[++reference_buf_idx_ % TINY_ENC_REF_BUF_CNT];
    //printf("[%s.%d]===>:, ref_buf: %p, next: %p\n", __FUNCTION__, __LINE__, ref_buf, next_ref_buf); 
    partition_encode(ref_buf, src,  y_coordinate_buf, out_frame, next_ref_buf, width, height, Y_BLOCK_WIDTH); //encode y
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    partition_encode(ref_buf + y_size, src + y_size,  uv_coordinate_buf, out_frame + y_size, next_ref_buf + y_size,  width / 2, height, UV_BLOCK_WIDTH); //encode uv
#else
    uint8_t* predictive_frame = new uint8_t[frame_size];
    memset(predictive_frame, 0, frame_size);
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    partition_encode2(src, reference_frame_buf_, predictive_frame, out_frame, coordinate_buf, width, height, Y_BLOCK_WIDTH, Y_BLOCK_SIZE); //encode y
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    //partition_encode2(src + y_size, reference_frame_buf_ + y_size, predictive_frame + y_size, out_frame + y_size, coordinate_buf + y_block_coordinate_size, width / 2, height, UV_BLOCK_WIDTH, UV_BLOCK_SIZE); 
#endif

#if 0
    uint16_t* coordinate_buf1 = (uint16_t*)coordinate_buf;
    for (size_t i = 0; i < block_coordinate_size / 4; i++)
    {
        uint16_t pos_x = *coordinate_buf1++;
        uint16_t pos_y = *coordinate_buf1++;        
        printf("[%s.%d]===>:posx: %d, posy:%d\n", __FUNCTION__, __LINE__, pos_x, pos_y); 
    }
#endif

#if 0
    static uint32_t idx = 0;

    if (idx++ >= 1)
    {
        dump_frame("ref_frame.nv12", reference_frame_buf_, frame_size);
        dump_frame("out_frame.nv12", out_frame, frame_size);
        //dump_frame("predictive_frame.nv12", predictive_frame, frame_size);
        abort();
    }    
#endif

#endif

    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    uint32_t out_buf_size = LZ4_COMPRESSBOUND(before_size);

    out_buf.buf = new uint8_t[out_buf_size];
    out_buf.release = [](KfaBuffer& buf) 
    {
        if (buf.buf)
        {
            delete[] buf.buf;
        }
    };
#if 1
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    int res_size = LZ4_compress_HC((const char*)out_frame_buf_, (char*)out_buf.buf, before_size, out_buf_size, LZ4HC_CLEVEL_MAX);
    
    if(res_size <= 0)
    {
        return false;
    } 
#else
    int res_size = ZSTD_compress((char*)out_buf.buf, out_buf_size, (const char*)out_frame_buf_, before_size, 22);
    if(res_size <= 0)
    {
        return false;
    }     
#endif

    //memcpy(reference_frame_buf_, src, src_size);

    printf("[%s.%d]===>: before: %d , after: %d\n", __FUNCTION__, __LINE__, before_size, res_size);
    out_buf.size = res_size;
    return true;
}






TinyEncoder2::TinyEncoder2()
{

}

TinyEncoder2::~TinyEncoder2()
{
    if (reference_frame_buf_)
    {
        delete[] reference_frame_buf_;
    }

    if (out_frame_buf_)
    {
        delete[] out_frame_buf_;
    }
}

bool TinyEncoder2::encode(const uint8_t* src, uint32_t src_size, KfaBuffer& out_buf) 
{
    if (!reference_frame_buf_)
    {
        //mean first frame, the only 'I' frame
        reference_frame_buf_ = new uint8_t[src_size];
        memcpy(reference_frame_buf_, src, src_size);

        out_frame_buf_ = new uint8_t[src_size];

        //reference_y_size_ = src_size;
        uint32_t out_buf_size = LZ4_COMPRESSBOUND(src_size);
        //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
        out_buf.buf = new uint8_t[out_buf_size];
        out_buf.release = [](KfaBuffer& buf) 
        {
            if (buf.buf)
            {
                delete[] buf.buf;
            }
        };

        int res_size = LZ4_compress_HC((const char*)src, (char*)out_buf.buf, src_size, out_buf_size, LZ4HC_CLEVEL_MAX);

        if(res_size <= 0)
        {
            return false;
        } 
        //printf("[%s.%d]===>: res_size:\n", __FUNCTION__, __LINE__);
        out_buf.size = res_size;
        return true;
    }

    uint32_t width = 1920; //todo fix me??
    uint32_t height = 720;
    uint32_t y_size = width * height;
    uint32_t uv_size = y_size / 2;
    uint32_t frame_size = y_size + uv_size; // only for nv12



   uint32_t out_buf_size = LZ4_COMPRESSBOUND(frame_size);

    out_buf.buf = new uint8_t[out_buf_size];
    out_buf.release = [](KfaBuffer& buf) 
    {
        if (buf.buf)
        {
            delete[] buf.buf;
        }
    };
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    int res_size = LZ4_compress_HC((const char*)out_frame_buf_, (char*)out_buf.buf, y_size, out_buf_size, LZ4HC_CLEVEL_MAX);
    
    if(res_size <= 0)
    {
        return false;
    } 

    memcpy(reference_frame_buf_, src, src_size);

    printf("[%s.%d]===>: before: %d , after: %d\n", __FUNCTION__, __LINE__, frame_size, res_size);
    out_buf.size = res_size;
    return true;

}